import os, sys
import pandas as pd
from PIL import Image

# 读取配置文件
from code.facenet import Facenet
from code.utils.utils import getRootPath


shape = [160, 160, 3]



def main(to_pred_dir, result_save_path, model):
    subdirs = os.listdir(to_pred_dir) # name
    subdirs = sorted(subdirs,key=lambda x:int(x))
    labels = []
    for subdir in subdirs:
        img1 = Image.open(os.path.join(to_pred_dir,subdir,"a.jpg"))
        img2 = Image.open(os.path.join(to_pred_dir,subdir,"b.jpg"))
        result = model.similar_probability(img1, img2)
        labels.append(1 if result<0.62 else 0)
        # labels.append(float(result))
        # print(float(result))

    # 字典中的key值即为csv中列名
    dataframe = pd.DataFrame({'id': subdirs, 'label': labels})
    # 将DataFrame存储为csv,index表示是否显示行名，default=True
    dataframe.to_csv(result_save_path, index=False, sep=',')



if __name__ == "__main__":
    rootpath = getRootPath()
    model_path = r'code/model/ep080-loss0.008401.pth'
    model_path = os.path.join(rootpath, model_path)
    model = Facenet(model_path=model_path)  # 加载模型
    to_pred_dir = sys.argv[1]
    result_save_path = sys.argv[2]
    # to_pred_dir = r'E:/Competition/dataset/train/data'
    # result_save_path = r'./result/result.csv'
    main(to_pred_dir, result_save_path, model)

